import pandas as pd
import statsmodels.api as sm
import pingouin as pg
from pyprocessmacro import Process
from utils.analysis.common.mediation_common_class import RegressionDetails, RegressionModel, Model14Result


def cal_f_value(cur_df: pd.DataFrame, x_cols: [], y_col: str):
    f_df = cur_df
    # 计算F值
    cs_name = '常量'  # 初始化常量
    f_df[cs_name] = 1
    # 提取自变量和因变量, 自变量，中介变量和控制变量合并为自变量
    f_x = f_df[x_cols]
    f_y = f_df[y_col]
    # 添加截距项
    f_x = sm.add_constant(f_x)
    # 创建多元线性回归模型
    f_result = sm.OLS(f_y, f_x).fit()
    cur_res = dict()
    cur_res['f'] = f_result.fvalue
    cur_res['p'] = f_result.f_pvalue
    return cur_res


def analysis_model8(cur_df: pd.DataFrame, x, y, m, w, covar):
    # 计算两列的均值
    mean_x = cur_df[x].mean()
    mean_w = cur_df[w].mean()
    # 计算新列的值
    mult_col_name = 'INT'
    cur_df[mult_col_name] = (cur_df[x] - mean_x) * (cur_df[w] - mean_w)

    x2 = [x, w, mult_col_name]
    x2.extend(covar)
    m2xw = pg.linear_regression(cur_df[x2], cur_df[m])
    m2xw_fp = cal_f_value(cur_df, x2, m)

    # x3需要的数据和x1一样
    x3 = [x, w, m, mult_col_name]
    x3.extend(covar)
    y2xwm = pg.linear_regression(cur_df[x3], cur_df[y])
    y2xwm_fp = cal_f_value(cur_df, x3, y)

    cur_res = dict()
    cur_res['regression_m2xw'] = m2xw  # m，x对y的回归
    cur_res['regression_m2xw_fp'] = m2xw_fp  # m，x对y的回归
    cur_res['regression_y2xm2'] = y2xwm  # x对m的回归
    cur_res['regression_y2xm2_fp'] = y2xwm_fp  # x对m的回归
    return cur_res


def do_process(cur_df: pd.DataFrame, x, y, m, w, controls):
    p = Process(data=cur_df, model=8, x=x, y=y, m=[m], w=w, controls=controls, center=True, controls_in="all")
    # p.summary()
    print(p.direct_model)
    print(p.indirect_model)


class MediationModel8:
    def __init__(self, df: pd.DataFrame, x_arr: [], y_arr: [], m_arr: [], w_arr: [], covar: []):
        self.df = df
        self.x_arr = x_arr
        self.y_arr = y_arr
        self.m_arr = m_arr
        self.w_arr = w_arr
        self.covar = covar

    def item_analysis(self, x, y, m, w):
        cur_df = self.df
        # 计算两列的均值
        mean_x = cur_df[x].mean()
        mean_w = cur_df[w].mean()
        # 计算新列的值
        mult_col_name = 'INT'
        cur_df[mult_col_name] = (cur_df[x] - mean_x) * (cur_df[w] - mean_w)

        x2 = [x, w, mult_col_name]
        x2.extend(self.covar)
        m2xw = pg.linear_regression(cur_df[x2], cur_df[m])
        m2xw_fp = cal_f_value(cur_df, x2, m)

        # x3需要的数据和x1一样
        x3 = [x, w, m, mult_col_name]
        x3.extend(self.covar)
        y2xwm = pg.linear_regression(cur_df[x3], cur_df[y])
        y2xwm_fp = cal_f_value(cur_df, x3, y)

        # 模型1
        model_y2xwm = RegressionModel(f=m2xw_fp.get('f'), p=m2xw_fp.get('p'),
                                      details=[RegressionDetails(*row) for row in m2xw.to_numpy()])
        # 模型2
        model_m2xw = RegressionModel(f=y2xwm_fp.get('f'), p=y2xwm_fp.get('p'),
                                     details=[RegressionDetails(*row) for row in y2xwm.to_numpy()])
        # 总返回结果
        res_obj = Model14Result(y2xwm=model_y2xwm, m2xw=model_m2xw, x=x, y=y, m=m, w=w, covar=self.covar)
        return res_obj

    def do_process(self, x, y, m, w):
        controls = self.covar
        p = Process(data=self.df, model=8, x=x, y=y, m=[m], w=w, controls=controls, center=True, controls_in="all")
        # p.summary()
        print(p.direct_model)
        print(p.indirect_model)
        return p

    def analysis(self):
        res = []
        for x in self.x_arr:
            for y in self.y_arr:
                for m in self.m_arr:
                    for w in self.w_arr:
                        item_res = self.item_analysis(x, y, m, w)
                        item_res.p = self.do_process(x, y, m, w)
                        res.append(item_res)
        return res


if __name__ == '__main__':
    pd.set_option('expand_frame_repr', False)
    df = pd.read_excel('./test_datas.xlsx')
    df = df[["性别", "年龄", "学历", "企业角色", 'JG', 'YY', 'JX', 'XW']]
    # 1.此文件的结果不包含整体的F值和p值
    # 2.若有多个x,y,m，则需要遍历进行获取结果
    # x_arr = ['JG']
    # y_arr = ['YY']
    # m_arr = ['XW']
    # w_arr = ['JX']
    # covar = ["性别", "年龄", "学历", "企业角色"]
    # res = []
    # for x in x_arr:
    #     for y in y_arr:
    #         for m in m_arr:
    #             for w in w_arr:
    #                 # 1. 线性回归
    #                 item_res = analysis_model8(df, x, y, m, w, covar)
    #                 res.append(item_res)
    #                 for cc in item_res:
    #                     print(cc)
    #                     print(item_res[cc])
    #                 print('========================================================================================')
    #                 # 2. process调节效应分析
    #                 do_process(df, x, y, m, w, covar)
    print('===================================================================================================')
